import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os


class ImageProcessor:
    def __init__(self, image_path):
        self.original_image = np.array(Image.open(image_path))
        self.gray_image = self.rgb_to_gray(self.original_image)

    def rgb_to_gray(self, rgb_image):
        r, g, b = rgb_image[:, :, 0], rgb_image[:, :, 1], rgb_image[:, :, 2]
        gray = 0.299 * r + 0.587 * g + 0.114 * b
        return gray.astype(np.uint8)

    def convolution2d(self, image, kernel):
        img_h, img_w = image.shape
        kernel_h, kernel_w = kernel.shape
        pad_h = kernel_h // 2
        pad_w = kernel_w // 2
        padded_image = np.zeros((img_h + 2 * pad_h, img_w + 2 * pad_w))
        padded_image[pad_h:pad_h + img_h, pad_w:pad_w + img_w] = image
        output = np.zeros_like(image, dtype=np.float32)
        for i in range(img_h):
            for j in range(img_w):
                output[i, j] = np.sum(
                    padded_image[i:i + kernel_h, j:j + kernel_w] * kernel
                )
        return output

    def sobel_filter(self):
        sobel_x = np.array([[-1, 0, 1],
                            [-2, 0, 2],
                            [-1, 0, 1]])
        sobel_y = np.array([[-1, -2, -1],
                            [0, 0, 0],
                            [1, 2, 1]])
        grad_x = self.convolution2d(self.gray_image, sobel_x)
        grad_y = self.convolution2d(self.gray_image, sobel_y)
        gradient_magnitude = np.sqrt(grad_x ** 2 + grad_y ** 2)
        gradient_magnitude = (gradient_magnitude / gradient_magnitude.max() * 255).astype(np.uint8)
        return gradient_magnitude

    def custom_filter(self, kernel):
        filtered = self.convolution2d(self.gray_image, kernel)
        filtered = (filtered - filtered.min()) / (filtered.max() - filtered.min()) * 255
        return filtered.astype(np.uint8)

    def color_histogram(self):
        histograms = np.zeros((256, 3))
        for channel in range(3):
            channel_data = self.original_image[:, :, channel].ravel()
            for pixel_value in channel_data:
                histograms[pixel_value, channel] += 1
        return histograms

    def extract_texture_features(self):
        features = {}
        img_normalized = self.gray_image / 255.0
        features['mean'] = np.mean(img_normalized)
        features['std'] = np.std(img_normalized)
        energy = np.sum(img_normalized ** 2)
        features['energy'] = energy / (self.gray_image.shape[0] * self.gray_image.shape[1])
        histogram = np.histogram(self.gray_image, bins=256, range=(0, 256))[0]
        probability = histogram / np.sum(histogram)
        entropy = -np.sum(probability * np.log2(probability + 1e-10))
        features['entropy'] = entropy / np.log2(256)
        local_variance = np.std(img_normalized) ** 2
        features['contrast'] = local_variance
        return features


def main():
    st.title("何明钊：机器视觉作业一")
    st.write("上传一张图片，展示Sobel滤波、自定义滤波、颜色直方图和纹理特征。")

    uploaded_file = st.file_uploader("上传图片", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image_path = os.path.join("temp", uploaded_file.name)
        os.makedirs("temp", exist_ok=True)
        with open(image_path, "wb") as f:
            f.write(uploaded_file.read())

        st.image(image_path, caption="原始图像", use_column_width=True)

        processor = ImageProcessor(image_path)

        # Sobel滤波
        sobel_result = processor.sobel_filter()
        st.image(sobel_result, caption="Sobel 滤波结果", use_column_width=True, clamp=True)

        # 自定义滤波
        custom_kernel = np.array([[1, 0, -1],
                                  [2, 0, -2],
                                  [1, 0, -1]])
        custom_result = processor.custom_filter(custom_kernel)
        st.image(custom_result, caption="自定义滤波结果", use_column_width=True, clamp=True)

        # 颜色直方图
        hist = processor.color_histogram()
        plt.figure(figsize=(10, 6))
        colors = ['red', 'green', 'blue']
        for i, color in enumerate(colors):
            plt.plot(hist[:, i], color=color, label=color.upper())
        plt.title('Color Histogram')
        plt.xlabel('Pixel Intensity')
        plt.ylabel('Frequency')
        plt.legend()
        st.pyplot(plt)

        # 纹理特征
        texture_features = processor.extract_texture_features()
        st.write("纹理特征：")
        st.json(texture_features)


if __name__ == "__main__":
    main()
